import time

import tensorflow as tf
print("TensorFlow 版本:", tf.__version__)

from tensorflow.keras import models
from tensorflow.keras.datasets import mnist
from tensorflow.keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPool2D
from tensorflow.keras.utils import to_categorical

gpu_list = tf.config.list_physical_devices('GPU')
print(f"gpu列表: {gpu_list}")

# 数据集的预处理
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(60000, 28, 28, 1) # 进行形状变换以适应卷积运算的需要
x_test = x_test.reshape(10000, 28, 28, 1)
y_train = to_categorical(y_train, 10) # 特征转换为one-hot编码,10位向量，代表数字0-9
y_test = to_categorical(y_test, 10) # 特征转换为one-hot编码

# CNN神经网络的添加和拟合
model = models.Sequential() # 序贯方式建立模型
model.add(Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)))  # 添加Conv2D二维卷积层
model.add(MaxPool2D(pool_size=(2,2)))
model.add(Conv2D(64, (3,3), activation='relu'))
model.add(MaxPool2D(pool_size=(2,2)))
model.add(Dropout(0.5))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(10, activation='softmax'))
model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy'])

start_time = time.time()
model.fit(x_train, y_train, validation_split=0.2, epochs=10, batch_size=200)
end_time = time.time()
print(f"训练耗时：{end_time - start_time}s")

# 利用测试集评估模型准确率
score = model.evaluate(x_test, y_test)
print("测试集预测准确率为：", score)

model.save('./cnn-hello-gpu.h5')
